import sys

def find_start_codes(data):
    """查找所有的起始码位置（0x00000001 或 0x000001）"""
    start_codes = []
    i = 0
    while i < len(data) - 4:
        if data[i:i+4] == b'\x00\x00\x00\x01':
            start_codes.append(i)
            i += 4
        elif data[i:i+3] == b'\x00\x00\x01':
            start_codes.append(i)
            i += 3
        else:
            i += 1
    return start_codes

def split_nalus(data):
    """根据起始码分割成 NALU 列表"""
    start_codes = find_start_codes(data)
    nalus = []
    for i in range(len(start_codes)):
        start = start_codes[i]
        end = start_codes[i + 1] if i + 1 < len(start_codes) else len(data)
        nalus.append(data[start:end])
    return nalus

def get_nalu_type(nalu):
    """从NALU中提取类型值"""
    if nalu.startswith(b'\x00\x00\x00\x01'):
        header = nalu[4]
    elif nalu.startswith(b'\x00\x00\x01'):
        header = nalu[3]
    else:
        return -1
    return header & 0x1F

def nalu_type_name(nalu_type):
    """返回 NALU 类型名称"""
    return {
        1: 'P帧',
        5: 'I帧',
        6: 'SEI',
        7: 'SPS',
        8: 'PPS',
        9: 'AUD',
    }.get(nalu_type, '未知')

def main(filepath):
    with open(filepath, 'rb') as f:
        data = f.read()

    nalus = split_nalus(data)
    grouped = {}

    for nalu in nalus:
        nalu_type = get_nalu_type(nalu)
        name = nalu_type_name(nalu_type)
        grouped.setdefault(name, []).append(nalu)

    for name in ['SPS', 'PPS', 'I帧', 'P帧', 'SEI', 'AUD', '未知']:
        if name in grouped:
            for idx, nalu in enumerate(grouped[name]):
                print(f"// {name} [{idx}] ")
                print(f"{list(nalu)},\n")

if __name__ == '__main__':
    main('output.h264')
